﻿using Resonance.TrainNode.Contracts.Entities;
using AForge.Neuro;
using System.Linq;
using System;
using AForge.Neuro.Learning;

namespace Resonance.TrainNode.Services
{
	internal static class JobWorkerThread
	{
		internal static void MainLoop(object jobDataObj)
		{
			JobInfo job = (JobInfo)jobDataObj;
			TrainJob tJob = job.Job;
			SessionInfo session = job.Session;
			TrainParameters trainParams = tJob.TrainParams;
			// TODO: don't have to access directly to session.Networks object ?
			NNet nnet = (from n in session.NNetworks where n.ID == tJob.NNetID select n).Single();

			tJob.StartTime = DateTime.Now;
			tJob.EndTime = null;

			try
			{
				// AForge initialization
				job.InputPatterns = GetInputPatterns(session);
				job.OutputPatterns = GetOutputPatterns(session);
				// create multi-layer neural network
				int[] layerNeuronCount = new int[nnet.Layers.Length];
				for (int l = 0; l < layerNeuronCount.Length; l++)
				{
					layerNeuronCount[l] = nnet.Layers[l].Neurons.Length;
				}
				ActivationNetwork aForgeNetwork = new ActivationNetwork(
					new BipolarSigmoidFunction(nnet.SigmoidAlpha),
					nnet.InputCount, layerNeuronCount);
				job.Network = aForgeNetwork;
				// create teacher
				BackPropagationLearning teacher = new BackPropagationLearning(job.Network);
				// set learning rate and momentum
				teacher.LearningRate = trainParams.LearningRate;
				teacher.Momentum = trainParams.Momentum;
				job.BackPropTeacher = teacher;

				tJob.CurrentEpoch = 0;
				tJob.CurrentError = double.MaxValue;
				double samples = (double)session.TrainPatterns.Count;
				bool workDone = false;

				while (tJob.State != JobStateEnum.Stopping && !workDone)
				{
					tJob.CurrentEpoch++;
					// TODO: could run one single pattern to stop the thread quickly in case of "Stopping" state
					tJob.CurrentError = teacher.RunEpoch(job.InputPatterns, job.OutputPatterns) / samples;

					if (tJob.CurrentEpoch >= trainParams.MaxEpochs ||
						tJob.CurrentError < trainParams.MaxError)
					{
						workDone = true;
					}
				}

				// free AForge objects
				job.InputPatterns = null;
				job.OutputPatterns = null;
				job.BackPropTeacher = null;

				tJob.EndTime = DateTime.Now;
				if (workDone)
					tJob.State = JobStateEnum.Completed;
				else
					tJob.State = JobStateEnum.Stopped;

				SessionManager.JobStateChangedCallback(job);
			}
			catch (Exception)
			{
				tJob.State = JobStateEnum.Error;
			}

			job.AssignedThread = null;
		}

		private static double[][] GetInputPatterns(SessionInfo si)
		{
			int samples = si.TrainPatterns.Count;
			double[][] patterns = new double[samples][];
			int neuronCount = si.TrainPatterns[0].InputValues.Length;
			for (int p = 0; p < samples; p++)
			{
				double[] neuronValues = new double[neuronCount];
				for (int v = 0; v < neuronCount; v++)
				{
					neuronValues[v] = si.TrainPatterns[p].InputValues[v];
				}
				patterns[p] = neuronValues;
			}
			return patterns;
		}

		private static double[][] GetOutputPatterns(SessionInfo si)
		{
			int samples = si.TrainPatterns.Count;
			double[][] patterns = new double[samples][];
			int neuronCount = si.TrainPatterns[0].OutputValues.Length;
			for (int p = 0; p < samples; p++)
			{
				double[] neuronValues = new double[neuronCount];
				for (int v = 0; v < neuronCount; v++)
				{
					neuronValues[v] = si.TrainPatterns[p].OutputValues[v];
				}
				patterns[p] = neuronValues;
			}
			return patterns;
		}
	}
}
